#include <iostream>
#include <algorithm>

// Included files for OpenGL Rendering
#ifdef __APPLE__
#include <OpenGL/gl.h>
#else
#include <GL/gl.h>
#endif

#include "photon_mapping.h"
#include "mesh.h"
#include "face.h"
#include "primitive.h"
#include "kdtree.h"
#include "utils.h"
#include "raytracer.h"
#include <math.h>
#include <map>
using namespace std;
// ========================================================================
// DESTRUCTOR
// ========================================================================
PhotonMapping::~PhotonMapping() {
  // cleanup all the photons
  delete kdtree;
}

// ========================================================================
// RECURSIVELY TRACE A SINGLE PHOTON
// ========================================================================
void PhotonMapping::TracePhoton(const Vec3f &position, const Vec3f &direction, const Vec3f &energy, int iter) 
{
	// ==============================================
	// ASSIGNMENT: IMPLEMENT RECURSIVE PHOTON TRACING
	// ==============================================
	
	Ray r=Ray(position,direction);
	Hit hit = Hit();
	if(raytracer->CastRay(r,hit,false)) // && iter < args->num_bounces)//&& energy > 0)
	{
		Vec3f spot=r.getOrigin()+(r.getDirection()*hit.getT());
		
		if(hit.getMaterial()->getReflectiveColor().Length() > 0) //if is reflective material
		{
			//MTRand MTwist = MTRand();
			//if(MTwist.rand(1) >= 0.5)
			//{
					Photon p=Photon(spot,direction,energy*hit.getMaterial()->getReflectiveColor(),iter);
					p.rEnergy=energy*hit.getMaterial()->getReflectiveColor();
					kdtree->AddPhoton(p);
				
				//cout<<"lol"<<endl;
				Vec3f normal = hit.getNormal();
				normal.Normalize();
				Vec3f reflect = direction - 2*(direction.Dot3(normal))*normal; 
				//reflect.Normalize();
				TracePhoton(spot,reflect,energy*hit.getMaterial()->getReflectiveColor(),iter+1);
			//}
		}
		else if(iter>0) //diffuse material otherwise
		{
			MTRand MTwist = MTRand();
			if(MTwist.rand(1) >= 0.5)
			{	
				Photon p=Photon(spot,direction,energy*hit.getMaterial()->getDiffuseColor(),iter);
				kdtree->AddPhoton(p);
				Vec3f randomDirection = RandomDiffuseDirection(hit.getNormal());
				TracePhoton(spot,randomDirection,energy*hit.getMaterial()->getDiffuseColor(),iter+1);
			}
			else if(iter > 0)
			{
				//kdtree->AddPhoton(p);
				//cout<<iter<<endl;
			}
		}
		
		
	}
	
	

}

// ========================================================================
// TRACE THE SPECIFIED NUMBER OF PHOTONS THROUGHOUT THE SCENE
// ========================================================================
void PhotonMapping::TracePhotons() {
  std::cout << "trace photons" << std::endl;

  // first, throw away any existing photons
  delete kdtree;

  // consruct a kdtree to store the photons
  BoundingBox *bb = mesh->getBoundingBox();
  Vec3f min = bb->getMin();
  Vec3f max = bb->getMax();
  Vec3f diff = max-min;
  min -= 0.001*diff;
  max += 0.001*diff;
  kdtree = new KDTree(BoundingBox(min,max));

  // photons emanate from the light sources
  const std::vector<Face*>& lights = mesh->getLights();

  // compute the total area of the lights
  double total_lights_area = 0;
  for (unsigned int i = 0; i < lights.size(); i++) {
    total_lights_area += lights[i]->getArea();
  }

  // shoot a constant number of photons per unit area of light source
  // (alternatively, this could be based on the total energy of each light)
  for (unsigned int i = 0; i < lights.size(); i++) {  
    double my_area = lights[i]->getArea();
    int num = args->num_photons_to_shoot * my_area / total_lights_area;
    // the initial energy for this photon
    Vec3f energy = my_area/double(num) * lights[i]->getMaterial()->getEmittedColor();
    Vec3f normal = lights[i]->computeNormal();
    for (int j = 0; j < num; j++) {
      Vec3f start = lights[i]->RandomPoint();
      // the initial direction for this photon (for diffuse light sources)
      Vec3f direction = RandomDiffuseDirection(normal);
      TracePhoton(start,direction,energy,0);
    }
  }
  
  //for(unsigned int i=0; i<kdtree->getPhotons().size(); i++)
  //{
//	kdtree->getPhotons()[i].setEnergy(kdtree->getPhotons()[i].getEnergy() * (1/kdtree->getPhotons.size()));
  //}
  
}


// ======================================================================
// PHOTON VISUALIZATION FOR DEBUGGING
// ======================================================================

void PhotonMapping::RenderPhotons() {
  RenderPhotonPositions();
  RenderPhotonDirections();
}

// render the position of each photon
void PhotonMapping::RenderPhotonPositions() {  
  if (kdtree == NULL) return;
  glDisable(GL_LIGHTING);
  glPointSize(3);
  glBegin(GL_POINTS);
  std::vector<const KDTree*> todo;  
  todo.push_back(kdtree);
  while (!todo.empty()) {
    const KDTree *node = todo.back();
    todo.pop_back(); 
    if (node->isLeaf()) {
      const std::vector<Photon> &photons = node->getPhotons();
      int num_photons = photons.size();
      for (int i = 0; i < num_photons; i++) {
	const Photon &p = photons[i];
	Vec3f energy = p.getEnergy()*args->num_photons_to_shoot;
	glColor3f(energy.x(),energy.y(),energy.z());
	const Vec3f &position = p.getPosition();
	glVertex3f(position.x(),position.y(),position.z());
      }
    } else {
      todo.push_back(node->getChild1());
      todo.push_back(node->getChild2());
    } 
  }
  glEnd();
  glEnable(GL_LIGHTING);
}


// render the incoming direction of each photon
void PhotonMapping::RenderPhotonDirections() {
  if (kdtree == NULL) return;
  glDisable(GL_LIGHTING);
  glLineWidth(1);
  glBegin(GL_LINES);
  std::vector<const KDTree*> todo;  
  todo.push_back(kdtree);
  BoundingBox *bb = mesh->getBoundingBox();
  double max_dim = bb->maxDim();
  while (!todo.empty()) {
    const KDTree *node = todo.back();
    todo.pop_back(); 
    if (node->isLeaf()) {
      const std::vector<Photon> &photons = node->getPhotons();
      int num_photons = photons.size();
      for (int i = 0; i < num_photons; i++) {
	const Photon &p = photons[i];
	const Vec3f a = p.getPosition();
	Vec3f b = p.getPosition()-(p.getDirectionFrom()*0.02*max_dim);
	Vec3f energy = p.getEnergy()*args->num_photons_to_shoot;
	glColor3f(energy.x(),energy.y(),energy.z());
	glVertex3f(a.x(),a.y(),a.z());
	glVertex3f(b.x(),b.y(),b.z());
      }
    } else {
      todo.push_back(node->getChild1());
      todo.push_back(node->getChild2());
    } 
  }
  glEnd();
  glEnable(GL_LIGHTING);
}


// ======================================================================
// render a simple wireframe of the KD tree
void PhotonMapping::RenderKDTree() {
  if (kdtree == NULL) return;
  glDisable(GL_LIGHTING);
  glLineWidth(1);
  glColor3f(0,0,0);
  glBegin(GL_LINES);
  std::vector<const KDTree*> todo;  
  todo.push_back(kdtree);
  while (!todo.empty()) {
    const KDTree *node = todo.back();
    todo.pop_back(); 
    if (node->isLeaf()) {
      const Vec3f& min = node->getMin();
      const Vec3f& max = node->getMax();
      
      glVertex3f(min.x(),min.y(),min.z());
      glVertex3f(max.x(),min.y(),min.z());
      glVertex3f(min.x(),min.y(),min.z());
      glVertex3f(min.x(),max.y(),min.z());
      glVertex3f(max.x(),max.y(),min.z());
      glVertex3f(max.x(),min.y(),min.z());
      glVertex3f(max.x(),max.y(),min.z());
      glVertex3f(min.x(),max.y(),min.z());

      glVertex3f(min.x(),min.y(),min.z());
      glVertex3f(min.x(),min.y(),max.z());
      glVertex3f(min.x(),max.y(),min.z());
      glVertex3f(min.x(),max.y(),max.z());
      glVertex3f(max.x(),min.y(),min.z());
      glVertex3f(max.x(),min.y(),max.z());
      glVertex3f(max.x(),max.y(),min.z());
      glVertex3f(max.x(),max.y(),max.z());
      
      glVertex3f(min.x(),min.y(),max.z());
      glVertex3f(max.x(),min.y(),max.z());
      glVertex3f(min.x(),min.y(),max.z());
      glVertex3f(min.x(),max.y(),max.z());
      glVertex3f(max.x(),max.y(),max.z());
      glVertex3f(max.x(),min.y(),max.z());
      glVertex3f(max.x(),max.y(),max.z());
      glVertex3f(min.x(),max.y(),max.z());
      
    } else {
      todo.push_back(node->getChild1());
      todo.push_back(node->getChild2());
    } 
  }
  glEnd();
  glEnable(GL_LIGHTING);
}


// ======================================================================

// helper function
bool closest_photon(const std::pair<Photon,double> &a, const std::pair<Photon,double> &b) {
  return (a.second < b.second);
}


// ======================================================================
Vec3f PhotonMapping::GatherIndirect(const Vec3f &point, const Vec3f &normal, const Vec3f &direction_from) const 
{
	if (kdtree == NULL) 
	{ 
		std::cout << "WARNING: Photons have not been traced throughout the scene." << std::endl;
		return Vec3f(0,0,0); 
	}
	
	KDTree* spotTree=kdtree->TraverseTree(point);
	Vec3f min=point-Vec3f(.1,.1,.1);
	Vec3f max=point+Vec3f(.1,.1,.1);

	BoundingBox box(min,max);
	BoundingBox box2(point-Vec3f(.5,.5,.5),point-Vec3f(.1,.1,.1));
	BoundingBox box3(point+Vec3f(.1,.1,.1),point+Vec3f(.5,.5,.5));
	BoundingBox box4(point+Vec3f(-.5,.1,.1),point+Vec3f(-.1,.5,.1));
	BoundingBox box5(point+Vec3f(.1,-.5,.1),point+Vec3f(.5,-.1,.1));
	
	
	vector<Photon> indirectPhotons;
	vector<Photon> indirectPhotons2;
	vector<Photon> indirectPhotons3;
	vector<Photon> indirectPhotons4;
	vector<Photon> indirectPhotons5;
	vector<vector<Photon> > vecs;
	vecs.push_back(indirectPhotons);
	
	vector<Photon> finalPhotons;
	vector<Vec3f> photonDirections;
	kdtree->CollectPhotonsInBox(box,indirectPhotons);
	kdtree->CollectPhotonsInBox(box2,indirectPhotons2);
	kdtree->CollectPhotonsInBox(box3,indirectPhotons3);
	kdtree->CollectPhotonsInBox(box4,indirectPhotons4);
	kdtree->CollectPhotonsInBox(box5,indirectPhotons5);
	
	
	bool different=false;
	
	for(unsigned int i=0; i<vecs.size(); i++)
	{
		if(indirectPhotons2.size() != vecs[i].size())
		{
			different=true;
		}
	}
	if(different) vecs.push_back(indirectPhotons2);
	different=false;
	
	for(unsigned int i=0; i<vecs.size(); i++)
	{
		if(indirectPhotons3.size() != vecs[i].size())
		{
			different=true;
		}
	}
	if(different) vecs.push_back(indirectPhotons3);
	different=false;
	
	for(unsigned int i=0; i<vecs.size(); i++)
	{
		if(indirectPhotons4.size() != vecs[i].size())
		{
			different=true;
		}
	}
	if(different) vecs.push_back(indirectPhotons4);
	different=false;
	
	for(unsigned int i=0; i<vecs.size(); i++)
	{
		if(indirectPhotons5.size() != vecs[i].size())
		{
			different=true;
		}
	}
	if(different) vecs.push_back(indirectPhotons5);
	different=false;
	
	
	for(unsigned int i=1; i<vecs.size(); i++)
	{
		for(unsigned int j=0; j<vecs[i].size(); j++)
		{
			indirectPhotons.push_back(vecs[i][j]);
		}
	}
	
	
	
	//cout << "size of indirect = " << indirectPhotons.size() << endl;
	map<double,Photon> photonMap;

	for(unsigned int i=0; i<indirectPhotons.size(); i++)
	{
		Vec3f a=indirectPhotons[i].getPosition();//-point;
		double dist=sqrt( (a.x()-point.x())*(a.x()-point.x()) + (a.y()-point.y())*(a.y()-point.y()) + (a.z()-point.z())*(a.z()-point.z()) );
		
		Photon p=Photon(indirectPhotons[i].getPosition(), indirectPhotons[i].getDirectionFrom(), indirectPhotons[i].getEnergy(), 0);
		photonMap[dist]=p;//a.Length()]=p;//indirectPhotons[i];
		
	}
	
	int counter=0;
	int k=60;//indirectPhotons.size()/10;//50;
	Vec3f totaldirection = Vec3f(0,0,0);
	for( map<double,Photon>::iterator iter = photonMap.begin(); iter != photonMap.end() && counter<k; ++iter )
	{
		Photon p=Photon( (*iter).second.getPosition(), (*iter).second.getDirectionFrom(), (*iter).second.getEnergy(), 0);
		finalPhotons.push_back(p);    //(*iter).second); 
		photonDirections.push_back(p.getDirectionFrom());
		totaldirection += p.getDirectionFrom();
		counter++;
		if((*iter).first > .3) break;
	}
	
	//cout<<"Indirect Size: "<<indirectPhotons.size()<<endl;
	//cout<<"Final Size: "<<finalPhotons.size()<<endl;
	//spotTree->CollectPhotonsInBox(spotTree->bbox,indirectPhotons);//getPhotons();
	
	Vec3f indirectEnergy=Vec3f(0,0,0);
	
	for(unsigned int i=0; i<finalPhotons.size(); i++)
	{
		
	
		Vec3f reflect = finalPhotons[i].getDirectionFrom() - 2*(finalPhotons[i].getDirectionFrom().Dot3(normal))*normal; 
		Ray r=Ray(point,reflect);
		Hit hit=Hit();
		
		Vec3f dista = finalPhotons[i].getPosition() - point;
		indirectEnergy += finalPhotons[i].getEnergy()*(1+(1-dista.Length()))*10;
		
		
		
	}	

	// ================================================================
	// ASSIGNMENT: GATHER THE INDIRECT ILLUMINATION FROM THE PHOTON MAP
	// ================================================================

	// Fix this, return the contribution from indirect illumination
	return indirectEnergy;
}










